package cn.tworice.cryption.util;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom;
import java.util.Base64;

public class AesGcmUtil {
    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/GCM/NoPadding";
    // GCM 推荐 12 字节 IV，Tag 长度 128 位
    private static final int IV_SIZE = 12;
    private static final int TAG_BIT_LENGTH = 128;
    private static final int KEY_SIZE = 128; // 可改为 128 或 192，需配合 JCE Unlimited Strength

    private static final SecureRandom RANDOM = new SecureRandom();

    /**
     * 动态生成 AES 密钥，返回 Base64 编码
     */
    public static String generateKey() throws Exception {
        KeyGenerator keyGen = KeyGenerator.getInstance(ALGORITHM);
        keyGen.init(KEY_SIZE, RANDOM);
        SecretKey key = keyGen.generateKey();
        byte[] keyBytes = key.getEncoded();
        return Base64.getEncoder().encodeToString(keyBytes);
    }

    /**
     * AES-GCM 加密
     * @param base64Key   Base64 编码的 AES 密钥
     * @param plaintext   明文字符串（UTF-8）
     * @return            加密结果，格式为 Base64(IV ‖ 密文)，IV 在前
     */
    public static String encrypt(String base64Key, String plaintext) throws Exception {
        // 1. 解码并构造 SecretKey
        byte[] keyBytes = Base64.getDecoder().decode(base64Key);
        SecretKeySpec keySpec = new SecretKeySpec(keyBytes, ALGORITHM);

        // 2. 生成随机 IV
        byte[] iv = new byte[IV_SIZE];
        RANDOM.nextBytes(iv);
        GCMParameterSpec spec = new GCMParameterSpec(TAG_BIT_LENGTH, iv);

        // 3. 加密
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, spec);
        byte[] cipherBytes = cipher.doFinal(plaintext.getBytes("UTF-8"));

        // 4. 拼接 IV + 密文，并 Base64 返回
        byte[] ivAndCipher = new byte[iv.length + cipherBytes.length];
        System.arraycopy(iv, 0, ivAndCipher, 0, iv.length);
        System.arraycopy(cipherBytes, 0, ivAndCipher, iv.length, cipherBytes.length);
        return Base64.getEncoder().encodeToString(ivAndCipher);
    }

    /**
     * AES-GCM 解密
     * @param base64Key     Base64 编码的 AES 密钥
     * @param base64IvCipher Base64(IV ‖ 密文) 字符串
     * @return              解密后的明文（UTF-8）
     */
    public static String decrypt(String base64Key, String base64IvCipher) throws Exception {
        // 1. 解码
        byte[] keyBytes = Base64.getDecoder().decode(base64Key);
        byte[] ivAndCipher = Base64.getDecoder().decode(base64IvCipher);

        // 2. 分离 IV 与密文
        byte[] iv = new byte[IV_SIZE];
        byte[] cipherBytes = new byte[ivAndCipher.length - IV_SIZE];
        System.arraycopy(ivAndCipher, 0, iv, 0, IV_SIZE);
        System.arraycopy(ivAndCipher, IV_SIZE, cipherBytes, 0, cipherBytes.length);

        // 3. 构造 Key 和参数
        SecretKeySpec keySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        GCMParameterSpec spec = new GCMParameterSpec(TAG_BIT_LENGTH, iv);

        // 4. 解密
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, spec);
        byte[] plainBytes = cipher.doFinal(cipherBytes);
        return new String(plainBytes, "UTF-8");
    }

    // 测试示例
    public static void main(String[] args) throws Exception {
        // 生成密钥
        String base64Key = generateKey();
        System.out.println("AES Key (Base64): " + base64Key);

        // 待加密文本
        String message = "Hello AES-GCM!";
        // 加密
        String encrypted = encrypt(base64Key, message);
        System.out.println("Encrypted (Base64 iv+cipher): " + encrypted);

        // 解密
        String decrypted = decrypt(base64Key, encrypted);
        System.out.println("Decrypted: " + decrypted);
    }
}
